import os  # 导入操作系统接口模块
import matplotlib.pyplot as plt  # 导入matplotlib库用于绘图
import numpy as np  # 导入numpy库用于数值计算
import scipy.signal  # 导入scipy库中的signal模块用于信号处理

# 定义一个类，用于绘制损失曲线
class LossPlotter:
    def __init__(self, log_dir):
        # 初始化方法，设置日志目录路径
        self.log_dir = log_dir
        # 设置训练损失文件的路径
        self.train_loss_file = os.path.join(log_dir, 'f_scores_train.txt')
        # 设置验证损失文件的路径
        self.val_loss_file = os.path.join(log_dir, 'f_scores_val.txt')
        # 加载训练损失数据
        self.losses = self.load_losses(self.train_loss_file)
        # 加载验证损失数据
        self.val_loss = self.load_losses(self.val_loss_file)
        # 验证损失数据存在的标志
        self.val_loss_flag = len(self.val_loss) > 0

    def load_losses(self, file_path):
        # 加载损失数据的方法
        try:
            # 打开文件并读取每行数据，转换为浮点数
            with open(file_path, 'r') as file:
                losses = [float(line.strip()) for line in file.readlines()]
            return losses  # 返回损失数据列表
        except FileNotFoundError:
            return []  # 文件未找到时返回空列表

    def loss_plot(self):
        # 绘制损失曲线的方法
        iters = range(len(self.losses))  # 生成迭代次数的范围

        plt.figure()  # 创建一个新的绘图窗口
        # 绘制训练损失曲线，红色实线，线宽为2
        plt.plot(iters, self.losses, 'red', linewidth=2, label='train f_score')

        if self.val_loss_flag:
            # 如果存在验证损失数据，绘制验证损失曲线，珊瑚色实线，线宽为2
            plt.plot(iters, self.val_loss, 'coral', linewidth=2, label='val f_score')

        # try:
        #     # 根据损失数据的长度确定平滑曲线的参数
        #     if len(self.losses) < 25:
        #         num = 5  # 数据少时，窗口大小为5
        #     else:
        #         num = 15  # 数据多时，窗口大小为15
        #
        #     # 绘制平滑后的训练损失曲线，绿色虚线，线宽为2
        #     plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle='--', linewidth=2,
        #              label='smooth train f_score')
        #
        #     if self.val_loss_flag:
        #         # 绘制平滑后的验证损失曲线，棕色虚线，线宽为2
        #         plt.plot(iters, scipy.signal.savgol_filter(self.val_loss, num, 3), '#8B4513', linestyle='--',
        #                  linewidth=2, label='smooth val f_score')
        # except Exception as e:
        #     # 平滑曲线时的异常处理，打印异常信息
        #     print(f"Exception during smoothing: {e}")

        plt.grid(True)  # 显示网格
        plt.xlabel('Epoch')  # 设置X轴标签为'Epoch'
        plt.ylabel('f_score')  # 设置Y轴标签为'f_score'
        plt.legend(loc="lower right")  # 显示图例并设置位置为右下角

        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)  # 如果日志目录不存在，则创建

        # 保存损失曲线图像到日志目录
        plt.savefig(os.path.join(self.log_dir, "1.png"))

        plt.cla()  # 清除当前活动的轴
        plt.close("all")  # 关闭所有图窗口

# 主函数
if __name__ == "__main__":
    # 定义存储损失数据和图片的文件夹路径
    log_directory = r"E:\postgraduate\大论文\seg_train\短数据集 500轮次"

    # 创建 LossPlotter 实例
    plotter = LossPlotter(log_directory)

    # 绘制损失曲线并保存图像
    plotter.loss_plot()
